import argparse
import asyncio
import fractions
import json
import os
import time

import cv2
import numpy as np
import pyautogui
from aiohttp import web
from aiortc import RTCPeerConnection, RTCSessionDescription, MediaStreamTrack
from aiortc.contrib.media import MediaRelay, MediaPlayer
from av import VideoFrame

ROOT = os.path.dirname(__file__)

relay = None
webcam = None


class VideoStreamTrack(MediaStreamTrack):
    kind = "video"

    def __init__(self):
        super().__init__()  # 初始化父类
        self._frame = None
        self._start_time = time.time()

    async def recv(self):
        # 假设我们想要30fps的帧率
        fps = 10
        time_base = 1 / fps

        # 计算当前帧的 pts
        current_time = time.time()
        elapsed_time = current_time - self._start_time
        pts = int(elapsed_time * fps)

        # 在这里捕获桌面截图并转换为 VideoFrame
        img = pyautogui.screenshot()
        frame = np.array(img)[:, :, :3]  # 删除alpha通道
        # 将 RGB 格式转换为 BGR 格式
        frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
        video_frame = VideoFrame.from_ndarray(frame, format='bgr24')
        video_frame.pts = pts
        video_frame.time_base = fractions.Fraction(1, fps)

        await asyncio.sleep(time_base)  # 控制帧率
        return video_frame


async def index(request):
    content = open(os.path.join(ROOT, "index.html"), "r").read()
    return web.Response(content_type="text/html", text=content)


async def javascript(request):
    content = open(os.path.join(ROOT, "client.js"), "r").read()
    return web.Response(content_type="application/javascript", text=content)


async def offer(request):
    params = await request.json()
    offer = RTCSessionDescription(sdp=params["sdp"], type=params["type"])

    pc = RTCPeerConnection()
    pcs.add(pc)

    @pc.on("connectionstatechange")
    async def on_connectionstatechange():
        print("Connection state is %s" % pc.connectionState)
        if pc.connectionState == "failed":
            await pc.close()
            pcs.discard(pc)

    # 创建视频轨道实例
    # 摄像头取流测试
    # options = {"framerate": "30", "video_size": "640x480"}
    # webcam = MediaPlayer("default:none", format="avfoundation", options=options)
    # video_track = webcam.video
    # 桌面截图
    video_track = VideoStreamTrack()

    # 使用 MediaRelay 订阅这个轨道以支持多个消费者
    relay = MediaRelay()
    video_relay_track = relay.subscribe(video_track)

    # 将 relay 订阅的轨道添加到 RTCPeerConnection 中
    pc.addTrack(video_relay_track)

    await pc.setRemoteDescription(offer)

    answer = await pc.createAnswer()
    await pc.setLocalDescription(answer)

    return web.Response(
        content_type="application/json",
        text=json.dumps(
            {"sdp": pc.localDescription.sdp, "type": pc.localDescription.type}
        ),
    )


pcs = set()


async def on_shutdown(app):
    # close peer connections
    coros = [pc.close() for pc in pcs]
    await asyncio.gather(*coros)
    pcs.clear()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="WebRTC webcam demo")
    parser.add_argument(
        "--host", default="0.0.0.0", help="Host for HTTP server (default: 0.0.0.0)"
    )
    parser.add_argument(
        "--port", type=int, default=8080, help="Port for HTTP server (default: 8080)"
    )
    args = parser.parse_args()

    app = web.Application()
    app.on_shutdown.append(on_shutdown)
    app.router.add_get("/", index)
    app.router.add_get("/client.js", javascript)
    app.router.add_post("/offer", offer)
    web.run_app(app, host=args.host, port=args.port)

